﻿using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Numerics;
using System.Text;
using System.Threading.Tasks;

namespace BTalk.Crypto.Algorithms
{
    public class RsaKey
    {
        public BigInteger N { get; set; }
        public BigInteger E { get; set; }
        public BigInteger D { get; set; }
    }

    public enum RsaMode
    {
        Normal,
        Hybrid
    }

    public class Rsa : ICryptoStrategy
    {
        private static Random random = new Random();
        private const int bytesInBlock = 4;

        private Tuple<RsaKey, RsaKey> _key;

        public dynamic Configuration
        {
            get
            {
                return _key;
            }
            set
            {
                if (!(value is Tuple<RsaKey, RsaKey>))
                    throw new ArgumentException();

                _key = value;
            }
        }
        public RsaMode Mode { get; set; }
        public ICryptoStrategy SymmetricAlgorithm { get; set; }

        public Rsa(RsaMode mode, ICryptoStrategy algorithm = null)
        {
            Mode = mode;
            SymmetricAlgorithm = algorithm;
        }

        public byte[] Encode(byte[] message)
        {

            if (Mode == RsaMode.Normal)
            {
                return EncodeRsa(message).SelectMany(m =>
                {
                    List<byte> list = new List<byte>();
                    list.AddRange(BitConverter.GetBytes(m.Length));
                    list.AddRange(m);
                    return list.ToArray();
                }).ToArray();
            }
            else
            {
                return  SymmetricAlgorithm.Encode(message);
            }
        }

        public byte[][] EncodeRsa(byte[] message)
        {
            IList<BigInteger> newValues = Pad(message).ToList();

            if (newValues.Any(v => Configuration.Item1.N <= v))
                throw new InvalidDataException();

            List<BigInteger> encrypted = newValues.AsParallel().AsOrdered().Select(v => BigInteger.ModPow(v, (BigInteger)Configuration.Item1.E, (BigInteger)Configuration.Item1.N)).ToList();

            return encrypted.Select(v => v.ToByteArray().Reverse().ToArray()).ToArray();
        }

        public byte[] DecodeRsa(byte[][] message)
        {
            List<BigInteger> integers = new List<BigInteger>();
            List<BigInteger> decrypted = new List<BigInteger>();

            foreach (var item in message)
            {
                integers.Add(new BigInteger(item.Reverse().ToArray()));
            }

            if (integers.Any(v => Configuration.Item2.N <= v))
                throw new InvalidDataException();

            decrypted = integers.AsParallel().AsOrdered().Select(i => (BigInteger)BigInteger.ModPow(i, (BigInteger)Configuration.Item2.D, (BigInteger)Configuration.Item2.N)).ToList();

            byte[] decryptedMessage = Unpad(decrypted.ToArray());

            return decryptedMessage;
        }

        private BigInteger[] Pad(byte[] message)
        {
            List<BigInteger> padding = new List<BigInteger>();
            var k = Configuration.Item1.N.ToByteArray().Length - 3;
            byte[] block = new byte[Configuration.Item1.N.ToByteArray().Length];
            BigInteger l = 0;
            block[0] = 0;
            block[1] = 2;
            int i = 2;
            if (k > message.Length)
            {
                byte[] filling = new byte[k - message.Length];

                foreach (var item in filling)
                {
                    block[i] = 255;
                    i++;
                }
            }
            block[i] = 0;
            i++;

            for (int d = 0; d < message.Length; d++)
            {
                if (i < block.Length)
                {
                    block[i] = message[d];
                }
                else
                {
                    padding.Add(new BigInteger(block.Reverse().ToArray()));
                    block = new byte[Configuration.Item1.N.ToByteArray().Length];
                    block[0] = 0;
                    block[1] = 2;
                    i = 2;
                    if (k > (message.Length - d))
                    {
                        byte[] filling = new byte[k - (message.Length - d)];

                        foreach (var item2 in filling)
                        {
                            block[i] = 255;
                            i++;
                        }
                    }
                    block[i] = 0;
                    i++;
                    block[i] = message[d];
                }
                i++;
            }

            padding.Add(new BigInteger(block.Reverse().ToArray()));

            return padding.ToArray();
        }

        private byte[] Unpad(BigInteger[] message)
        {
            List<byte> result = new List<byte>();
            foreach (var msg in message)
            {
                var data = msg.ToByteArray().Reverse().Skip(1).SkipWhile(d => d == 255).Skip(1).ToArray();
                result.AddRange(data);
            }
            return result.ToArray();
        }

        public byte[] Decode(byte[] message)
        {
            if (Mode == RsaMode.Normal)
            {

                BinaryReader reader = new BinaryReader(new MemoryStream(message));
                List<byte[]> numbers = new List<byte[]>();

                while (reader.PeekChar() != -1)
                {
                    int size = reader.ReadInt32();
                    numbers.Add(reader.ReadBytes(size));
                }
                reader.Close();

                return DecodeRsa(numbers.ToArray());
            }
            else
            {
                return SymmetricAlgorithm.Decode(message);
            }
        }
    }
}
